/**
 * 版权所有 2009-2012山东新北洋信息技术股份有限公司
 * 保留所有权利。
 */
package com.linyaonan.leetcode.medium._814;

/**
 * 给定二叉树根结点 root ，此外树的每个结点的值要么是 0，要么是 1。
 * <p>
 * 返回移除了所有不包含 1 的子树的原二叉树。
 * <p>
 * ( 节点 X 的子树为 X 本身，以及所有 X 的后代。)
 * <p>
 * 示例1:
 * 输入: [1,null,0,0,1]
 * 输出: [1,null,0,null,1]
 * <p>
 * 解释:
 * 只有红色节点满足条件“所有不包含 1 的子树”。
 * 右图为返回的答案。
 *
 * @ProjectName: leetcode
 * @Package: com.linyaonan.leetcode.medium._814
 * @ClassName: BinaryTreePruning
 * @Author: linyaonan
 * @Date: 2020/1/20 15:12
 */
public class BinaryTreePruning {

    public TreeNode pruneTree(TreeNode root) {
        TreeNode p = new TreeNode(-1);
        p.left = root;
        dfs(root, p, true);
        return p.left;
    }

    private void dfs(TreeNode node, TreeNode p, boolean isLeft) {
        if (node != null) {
            dfs(node.left, node, true);
            dfs(node.right, node, false);
            if (node.left == null && node.right == null && node.val == 0) {
                if (isLeft) {
                    p.left = null;
                } else {
                    p.right = null;
                }
            }
        }
    }
}
